from .base import HookBase
from .priority import Priority


class SummaryHook(HookBase):
    def __init__(self, interval=1, priority=Priority.LOWEST):
        super(SummaryHook, self).__init__(priority)
        self.interval = interval

    def after_epoch(self, runner, tag='train'):
        if not runner.is_master:
            return
        for k, v in runner.outputs.items():
            runner.tbwriter.add_scalar(f'{tag}/{k}', v, runner.epoch)

    def after_train_epoch(self, runner):
        self.after_epoch(runner, 'train')

    def after_val_epoch(self, runner):
        self.after_epoch(runner, 'val')
